#!/usr/bin/env python
"""
Experimental force-based algorithm for drawing a graph.
"""

from __future__ import division

import sys
import pygame
import random
import math


def add_t(*tuples):
    return tuple(sum(x) for x in zip(*tuples))

def mul_t(*tuples):
    def prod(x):
        n = 1
        for i in x:
            n *= i
        return n
    return tuple(prod(x) for x in zip(*tuples))

def mul_t_constant(t, constant):
    return tuple(x * constant for x in t)


class BadConfigFile(Exception) : pass

def read_config_file(filepath):
    with open(filepath) as configfile:
        raw_text = configfile.read().split('\n')
    config = {}
    for index, line in enumerate(raw_text):
        if line.startswith('#'):
            continue
        components = line.split(':')
        if len(components) > 2:
            raise BadConfigFile("Line " + str(index) + " - Invalid line")
        if len(components) < 2:
            continue
        config[components[0].strip()] = float(components[1].strip())
    return config



class Graph(object):
    def __init__(self, repulsion_const=1200000, attraction_const=2):
        self.repulsion = repulsion_const
        self.attraction = attraction_const
        self.timestep = 0.15
        self.damping = 0.5
        self.min_energy = 0.01
        self.min_spring_distance = 60
        self.num_nodes = 20
        self.max_neighbors = 4
        self.min_neighbors = 3
        return
    
    def generate_graph(self):
        nodes = self.generate_node_line()
        nodes = self.tangle(nodes)
        return nodes
    
    def generate_node_line(self):
        nodes_list = []
    
        n1 = Node()
        nodes_list.append(n1)
        nodes_left = self.num_nodes - 1
        while nodes_left:
            n2 = Node()
            n1.neighbors.append(n2)
            n2.neighbors.append(n1)
            n1 = n2
            nodes_list.append(n2)
            nodes_left -= 1
        return nodes_list
    
    def tangle(self, nodes_list):
        for node in nodes_list:
            size = self.max_neighbors - len(node.neighbors)
            min_size = self.min_neighbors - len(node.neighbors)
            if size < 1:
                continue
            print min_size, size-1
            num_neighbors = random.randint(0, size-1)
            while num_neighbors:
                other = node
                while other == node:
                    other = random.choice(nodes_list)
                    if (len(other.neighbors) > self.max_neighbors):
                        break
                node.neighbors.append(other)
                other.neighbors.append(node)
                num_neighbors -= 1
        print ''
        return nodes_list

    def load_params(self, parems):
        self.repulsion = parems['repulsion']
        self.attraction = parems['attraction']
        self.timestep = parems['timestep']
        self.damping = parems['damping']
        self.min_energy = parems['min_energy']
        self.min_spring_distance = parems['stable_spring_distance']
        self.num_nodes = parems['num_nodes']
        self.max_neighbors = parems['max_neighbors']
        self.min_neighbors = parems['min_neighbors']
        return
    
    def adjust_positions(self, nodes_list):
        kinetic_energy = 0
        for node in nodes_list:
            if node.static == True:
                continue
            net_force = (0,0)
                
            # repulsion to every other ring
            for other in nodes_list:
                if other != node:
                    nx, ny = net_force
                    cx, cy = self.coulomb_repulsion(node, other)
                    net_force = (nx + cx, ny + cy)
                
            # attraction based on string
            for other in node.neighbors:
                nx, ny = net_force
                hx, hy = self.hooke_attraction(node, other)
                net_force = (nx + hx, ny + hy)             

            # dampening
            temp = mul_t_constant(net_force, self.timestep)
            node.velocity = mul_t_constant(add_t(node.velocity, temp), self.damping)
            node.position = add_t(node.position, mul_t_constant(node.velocity, self.timestep))

            kinetic_energy += node.mass * (math.hypot(node.velocity[0], node.velocity[1]))
        return nodes_list, kinetic_energy

    def hooke_attraction(self, node, other):
        distance = self.find_distance(node, other)
        angle = self.angle(node, other)

        force = -self.attraction * (distance - self.min_spring_distance)

        x = math.cos(angle) * force
        y = math.sin(angle) * force
        return (x, y)

    def coulomb_repulsion(self, node, other):
        distance = self.find_distance(node, other)
        angle = self.angle(node, other)
        
        try:
            force = (self.repulsion * node.charge * other.charge) / (distance**2)
        except ZeroDivisionError:
            force = 0

        x = math.cos(angle) * force
        y = math.sin(angle) * force
        return (x, y)
        
    def delta(self, node, other):
        return (node.x - other.x, node.y - other.y)

    def find_distance(self, node, other):
        delta_x, delta_y = self.delta(node, other)
        distance = math.hypot(delta_x, delta_y)
        return distance

    def angle(self, node, other):
        delta_x, delta_y = self.delta(node, other)
        return math.atan2(delta_y, delta_x)
       


class Node(object):
    def __init__(self, position=None, neighbors=None, x_range=(0,800), y_range=(0,800), mass=1):
        self._position = (0,0)
        self.x_range = x_range
        self.y_range = y_range
        if not position:
            xmin, xmax = x_range
            ymin, ymax = y_range
            self.position = (random.randint(xmin, xmax), random.randint(ymin,ymax))
        else:
            self.position = position
        if not neighbors:
            self.neighbors = []
        else:
            self.neighbors = neighbors
        self.charge = 1
        self.velocity = (0,0)
        self.mass = mass
        self.static = False
        return
    
    def _get_position(self):
        return self._position
    def _set_position(self, value):
        self.x, self.y = value
        return
    position = property(_get_position, _set_position)
    
    def _get_x(self):
        x, y = self._position
        return x
    def _set_x(self, value):
        #xmin, xmax = self.x_range
        #if value < xmin + 15:
        #    value = xmin
        #if value > xmax - 15:
        #    value = xmax
        x, y = self._position
        self._position = (value, y)
        return
    x = property(_get_x, _set_x)
    
    def _get_y(self):
        x, y = self._position
        return y
    def _set_y(self, value):
        #ymin, ymax = self.y_range
        #if value < ymin + 15:
        #    value = ymin
        #if value > ymax - 15:
        #    value = ymax
        x, y = self._position
        self._position = (x, value)
        return
    y = property(_get_y, _set_y)

    def _get_rect(self):
        return pygame.Rect(self.x - 5, self.y - 5, 10, 10)
    rect = property(_get_rect)
    
        
    
    def __repr__(self):
        return str(self.position) + ': ' + str(self.neighbors)

class Graphics(object):
    def __init__(self, size=(800,800), caption="Nodes"):
        self.size = size
        self.caption = caption
        self.node_radius = 5
        self.node_color = (255,255,255)
        self.line_color = (255,255,255)
        self.screen_color = (0,0,0)
        
        pygame.init()
        self.screen = pygame.display.set_mode(self.size)
        pygame.display.set_caption(self.caption)
        return
    
    def draw_graph(self, nodes):
        self.screen.fill(self.screen_color)
        self.draw_lines(nodes)
        self.draw_nodes(nodes)
        return
        
    def draw_lines(self, nodes):
        for node in nodes:
            for neighbor in node.neighbors:
                pygame.draw.aaline(
                    self.screen,
                    self.line_color,
                    node.position,
                    neighbor.position,
                    1)
        return
    
    def draw_nodes(self, nodes):
        for node in nodes:
            intpos = (int(node.position[0]), int(node.position[1]))
            pygame.draw.circle(
                self.screen, 
                self.node_color, 
                intpos,
                self.node_radius,
                0)
        return
    
    def render_screen(self):
        pygame.display.flip()
        return
    
    def mainloop(self, node):
        is_running = True
        is_frozen = False
        is_node_selected = False
        selected_node = None
        
        grapher = Graph()
        grapher.load_params(read_config_file(r"nodesConfig.txt"))

        font = pygame.font.Font(None, 20)
        
        while is_running:
            # Events handling
            event = pygame.event.poll()
            if event.type == pygame.QUIT:
                pygame.quit()
                is_running = False
                sys.exit(0)
            if event.type == pygame.KEYUP:
                if event.key == pygame.K_n: # New nodes.
                    grapher.load_params(read_config_file(r"nodesConfig.txt"))
                    node = grapher.generate_graph()
                if event.key == pygame.K_f: # freeze.
                    is_frozen = not is_frozen
                if event.key == pygame.K_r: 
                    grapher.load_params(read_config_file(r"nodesConfig.txt"))
            if event.type == pygame.MOUSEBUTTONDOWN:
                if is_node_selected == False:
                    for n in node:
                        if n.rect.collidepoint(event.pos):
                            is_node_selected = True
                            selected_node = n
                            selected_node.static = True
                            break
            if event.type == pygame.MOUSEMOTION:
                if is_node_selected:
                    selected_node.position = event.pos
            if event.type == pygame.MOUSEBUTTONUP:
                if is_node_selected:
                    is_node_selected = False
                    selected_node.static = False
                    selected_node = None

            # Calculating and drawing.
            if is_frozen == False:
                node, kinetic_energy = grapher.adjust_positions(node)
                self.draw_graph(node)
                if kinetic_energy < 0.0001:
                    text = "Kinetic Energy: Practically zero"
                else:
                    text = "Kinetic Energy: " + str(kinetic_energy)
                f = font.render(text, 1, (255,255,255))
                self.screen.blit(f, (0,0))
                self.render_screen()
        return
    
def main():
    g = Graph()
    nodes = g.generate_graph()
    m = Graphics()
    m.draw_graph(nodes)
    m.mainloop(nodes)
    return

if __name__ == '__main__':
    main()
